import cv2
import numpy as np
from typing import Optional, Tuple, Dict, Any
import time
import threading
from queue import Queue, Empty, Full
from concurrent.futures import ThreadPoolExecutor, as_completed
import sys
import os
current_dir = os.path.dirname(os.path.abspath(__file__))
# 添加项目根目录到Python路径
sys.path.append(os.path.join(current_dir, '../'))  # 根据实际位置调整层级

# 假设这些模块已经存在
from image_provider import ImageProvider
from logger import setup_logger
from seg_inference2 import seg_func
from pose_inference import pose_func
from py_utils.rknn_executor import RKNN_model_container

# 初始化日志
logger = setup_logger()

segmodelPath = "./models/merge_seg_int8_rk3588.rknn"
posemodelPath = "./models/merge_pose_int8_rk3588.rknn"

class ImageData:
    """图像数据结构"""
    def __init__(self):
        self.timestamp: float = 0.0
        self.readable_timestamp: str = ""
        self.color_img: np.ndarray = None
        self.depth_img: np.ndarray = None
        self.color_path: Optional[str] = None

def initRKNN(rknnModel="./models/merge_yolov11_seg.rknn", id=0):
    """初始化单个RKNN模型"""
    rknn_lite = RKNN_model_container(rknnModel, "rk3588", id)
    return rknn_lite

def initRKNNs(rknnModel="./models/merge_yolov11_seg.rknn", TPEs=1, npus=[0]):
    """初始化多个RKNN模型"""
    rknn_list = []
    for npu in npus:
        for i in range(TPEs):
            rknn_list.append(initRKNN(rknnModel, npu))
    return rknn_list

class FixedRknnPoolExecutor:
    """修复的RKNN线程池执行器"""
    
    def __init__(self, rknnModel, TPEs, func, npus=[0]):
        self.TPEs = TPEs
        self.func = func
        self.npus = npus
        self.num = 0
        
        # 初始化RKNN模型池
        logger.info(f"初始化RKNN模型池: {rknnModel}, TPEs={TPEs}, NPUs={npus}")
        self.rknnPool = initRKNNs(rknnModel, TPEs, npus)
        
        # 创建线程池
        total_workers = TPEs * len(npus)
        self.pool = ThreadPoolExecutor(max_workers=total_workers)
        
        # 任务管理
        self.task_id = 0
        self.pending_futures = {}  # {task_id: future}
        self.completed_results = {}  # {task_id: (result, success)}
        self.lock = threading.Lock()
        
        logger.info(f"RKNN线程池初始化完成，总工作线程数: {total_workers}")

    def put(self, frame):
        """提交任务并返回任务ID"""
        with self.lock:
            task_id = self.task_id
            self.task_id += 1
        
        # 选择RKNN模型（轮询）
        model_index = self.num % len(self.rknnPool)
        self.num += 1
        
        # 添加调试信息
        logger.debug(f"提交RKNN任务 {task_id}: 使用模型索引={model_index}, 图像形状={frame.shape if hasattr(frame, 'shape') else 'unknown'}")
        
        # 提交任务 - 添加包装函数用于调试
        def debug_inference_wrapper(model, img):
            try:
                logger.debug(f"开始执行推理任务 {task_id}")
                start_time = time.time()
                result = self.func(model, img)
                end_time = time.time()
                logger.debug(f"推理任务 {task_id} 完成，耗时: {(end_time-start_time)*1000:.1f}ms, 结果: {type(result)}")
                if result is not None:
                    if isinstance(result, np.ndarray):
                        logger.debug(f"推理任务 {task_id} 结果数组形状: {result.shape}")
                    elif isinstance(result, (list, tuple)):
                        logger.debug(f"推理任务 {task_id} 结果列表长度: {len(result)}")
                    else:
                        logger.debug(f"推理任务 {task_id} 结果值: {result}")
                else:
                    logger.warning(f"推理任务 {task_id} 返回None!")
                return result
            except Exception as e:
                logger.error(f"推理任务 {task_id} 异常: {e}")
                raise
        
        future = self.pool.submit(debug_inference_wrapper, self.rknnPool[model_index], frame)
        
        # 保存任务
        with self.lock:
            self.pending_futures[task_id] = future
        
        # 添加完成回调
        future.add_done_callback(lambda f, tid=task_id: self._task_completed(tid, f))
        
        return task_id

    def _task_completed(self, task_id, future):
        """任务完成回调"""
        try:
            result,success = future.result()
            # print("result,sucess:", result, success)
        except Exception as e:
            logger.error(f"RKNN任务 {task_id} 执行失败: {e}")
            result = None
            success = False
        
        with self.lock:
            # 移动到完成队列
            self.completed_results[task_id] = (result, success)
            # 从待处理队列移除
            if task_id in self.pending_futures:
                del self.pending_futures[task_id]

    def get(self):
        """获取最早完成的结果"""
        with self.lock:
            if len(self.completed_results) == 0:
                # logger.info(f"没有完成的推理结果,睡眠 50 ms...")
                # time.sleep(0.05)
                return None, None, False
            
            # 获取最早的任务ID
            earliest_task_id = min(self.completed_results.keys())
            result, success = self.completed_results.pop(earliest_task_id)
            # print("result,success:", result, success)
            logger.info(f"获取RKNN任务结果 {earliest_task_id}: 结果类型: {type(result)}, 是否成功: {success}")
            return earliest_task_id, result, True

    def has_results(self):
        """检查是否有结果可获取"""
        with self.lock:
            return len(self.completed_results) > 0

    def pending_count(self):
        """获取待处理任务数量"""
        with self.lock:
            return len(self.pending_futures)

    def close(self):
        """关闭线程池"""
        logger.info("关闭RKNN线程池...")
        self.pool.shutdown(wait=True)
        for rknn_lite in self.rknnPool:
            rknn_lite.release()
        logger.info("RKNN线程池已关闭")

class AsyncProcessor:
    """异步处理器"""
    
    def __init__(self, seg_pool, pose_pool, target_fps=30, max_pending=50):
        self.seg_pool = seg_pool
        self.pose_pool = pose_pool
        self.target_fps = target_fps
        self.max_pending = max_pending
        self.running = False
        
        # 模拟图像提供者
        self.image_provider = ImageProvider(from_folder=False)
        
        # 任务管理
        self.task_queue = Queue(maxsize=max_pending)
        self.pending_tasks: Dict[int, Dict] = {}
        self.next_task_id = 0
        self.result_queue = Queue()
        
        # 帧率统计
        self.submit_count = 0
        self.process_count = 0
        self.last_log_time = time.time()
        
        # 性能统计
        self.total_inference_time = 0
        self.total_frames_processed = 0

    def start(self):
        """启动异步处理器"""
        self.running = True
        
        # 启动生产者线程
        self.producer_thread = threading.Thread(target=self._producer, name="Producer")
        self.producer_thread.daemon = True
        self.producer_thread.start()
        
        # 启动任务分发线程
        self.dispatch_thread = threading.Thread(target=self._dispatch_tasks, name="Dispatcher")
        self.dispatch_thread.daemon = True
        self.dispatch_thread.start()
        
        # 启动消费者线程
        self.consumer_thread = threading.Thread(target=self._consumer, name="Consumer")
        self.consumer_thread.daemon = True
        self.consumer_thread.start()
        
        logger.info("异步处理器启动完成")

    def stop(self):
        """停止异步处理器"""
        logger.info("正在停止异步处理器...")
        self.running = False
        
        # 等待线程结束
        for thread in [self.producer_thread, self.dispatch_thread, self.consumer_thread]:
            thread.join(timeout=2.0)
            if thread.is_alive():
                logger.warning(f"线程 {thread.name} 未能及时结束")
        
        # 清理线程池
        self.seg_pool.close()
        self.pose_pool.close()
        logger.info("异步处理器已停止")

    def _producer(self):
        """生产者线程：生成图像数据并提交任务"""
        frame_interval = 1.0 / self.target_fps
        
        while self.running:
            start_time = time.time()
            
            # 从图像源获取图像
            image_data = self.image_provider.get_next_image()
            if image_data is None:
                logger.info("图像源已结束")
                break
            
            # 创建任务
            task_id = self.next_task_id
            self.next_task_id += 1
            
            task = {
                'id': task_id,
                'image_data': image_data,
                'seg_id': None,
                'pose_id': None,
                'seg_result': None,
                'pose_result': None,
                'seg_success': False,
                'pose_success': False,
                'start_time': time.time(),
                'state': 'queued'
            }
            
            # 放入待处理任务
            self.pending_tasks[task_id] = task
            
            try:
                self.task_queue.put(task, timeout=0.1)
                self.submit_count += 1
                task['state'] = 'submitted'
            except Full:
                logger.warning(f"任务队列已满，任务 {task_id} 被丢弃")
                del self.pending_tasks[task_id]

    def _dispatch_tasks(self):
        """任务分发线程：将任务分发给RKNN线程池"""
        while self.running or not self.task_queue.empty():
            try:
                task = self.task_queue.get(timeout=0.5)
                
                # 提交到seg和pose线程池
                seg_id = self.seg_pool.put(task['image_data'].color_img)
                pose_id = self.pose_pool.put(task['image_data'].color_img)
                
                # 更新任务状态
                task['seg_id'] = seg_id
                task['pose_id'] = pose_id
                task['state'] = 'processing'
                
                logger.debug(f"任务 {task['id']} 已分发: seg_id={seg_id}, pose_id={pose_id}")
                
            except Empty:
                continue
            except Exception as e:
                logger.error(f"任务分发错误: {e}")

    def _consumer(self):
        """消费者线程：收集结果"""
        check_interval = 0.05  # 50ms检查一次
        
        while self.running or len(self.pending_tasks) > 0:
            results_found = False
            # 检查pose结果
            pose_id, pose_result, pose_success = self.pose_pool.get()
            # print("pose_result", pose_result, "pose_id", pose_id)
            if pose_id is not None:
                task = self._find_task_by_pose_id(pose_id)
                if task:
                    task['pose_result'] = pose_result
                    task['pose_success'] = True
                    results_found = True
                    logger.debug(f"收到pose结果: task_id={task['id']}, success={pose_success}")
            
            # 检查seg结果
            seg_id, seg_result, seg_success = self.seg_pool.get()
            if seg_id is not None:
                task = self._find_task_by_seg_id(seg_id)
                if task:
                    task['seg_result'] = seg_result
                    task['seg_success'] = True
                    results_found = True
                    logger.debug(f"收到seg结果: task_id={task['id']}, success={seg_success}")
            
            
            # 处理完成的任务
            self._process_completed_tasks()
            
            self._log_fps()
            # 如果没有结果，短暂休眠
            if not results_found:
                time.sleep(check_interval)
            
            # 更新统计

    def _find_task_by_seg_id(self, seg_id):
        """通过seg_id查找任务"""
        for task in self.pending_tasks.values():
            if task.get('seg_id') == seg_id:
                return task
        return None

    def _find_task_by_pose_id(self, pose_id):
        """通过pose_id查找任务"""
        for task in self.pending_tasks.values():
            if task.get('pose_id') == pose_id:
                return task
        return None

    def _process_completed_tasks(self):
        """处理已完成的任务"""
        completed_task_ids = []
        
        for task_id, task in list(self.pending_tasks.items()):
            # 检查是否两个结果都完成
            seg_done = task['seg_success'] 
            pose_done = task['pose_success']
            
            if seg_done and pose_done:
                logger.info(f"任务 {task_id} 已完成")
                # 计算推理时间
                inference_time = time.time() - task['start_time']
                self.total_inference_time += inference_time
                self.total_frames_processed += 1
                # print("task:",task)
                # 构建结果
                result = {
                    'task_id': task_id,
                    'image_data': task['image_data'],
                    'seg_result': task.get('seg_result'),
                    'pose_result': task.get('pose_result'),
                    'seg_success': task.get('seg_success'),
                    'pose_success': task.get('pose_success'),
                    'inference_time': inference_time,
                }
                
                # 放入结果队列
                try:
                    self.result_queue.put(result, timeout=0.1)
                    completed_task_ids.append(task_id)
                    self.process_count += 1
                    
                    logger.debug(f"任务 {task_id} 完成: "
                               f"seg_ok={result['seg_success']}, "
                               f"pose_ok={result['pose_success']}, "
                               f"time={inference_time*1000:.1f}ms")
                    
                except Full:
                    logger.warning("结果队列已满，丢弃结果")
                
                # 清理任务
                del self.pending_tasks[task_id]
        
        return len(completed_task_ids) > 0

    def _log_fps(self):
        """记录帧率信息"""
        current_time = time.time()
        if current_time - self.last_log_time >= 2.0:  # 每2秒输出一次
            elapsed = current_time - self.last_log_time
            
            # 计算帧率
            submit_fps = self.submit_count / elapsed if elapsed > 0 else 0
            process_fps = self.process_count / elapsed if elapsed > 0 else 0
            
            # 计算平均推理时间
            avg_inference = 0
            if self.total_frames_processed > 0:
                avg_inference = self.total_inference_time / self.total_frames_processed
            
            # 队列状态
            pending_tasks = len(self.pending_tasks)
            seg_pending = self.seg_pool.pending_count()
            pose_pending = self.pose_pool.pending_count()
            
            logger.info(
                f"性能统计: 提交={submit_fps:.1f}FPS, "
                f"处理={process_fps:.1f}FPS, "
                f"队列深度={pending_tasks}, "
                f"seg待处理={seg_pending}, "
                f"pose待处理={pose_pending}, "
                f"平均延迟={avg_inference*1000:.1f}ms "
                f"(目标={self.target_fps}FPS)"
            )
            
            # 重置计数器
            self.submit_count = 0
            self.process_count = 0
            self.last_log_time = current_time

    def get_result(self, block=False, timeout=None):
        """获取处理结果"""
        try:
            return self.result_queue.get(block=block, timeout=timeout)
        except Empty:
            return None

    def has_result(self):
        """检查是否有结果"""
        return not self.result_queue.empty()

    def pending_count(self):
        """获取待处理任务数量"""
        return len(self.pending_tasks)

# 初始化线程池
def initialize_pools():
    """初始化分割和姿态线程池"""
    seg_tpes = 6  # 每个NPU使用1个线程
    pose_tpes = 6
    
    logger.info(f"初始化分割RKNN线程池，每NPU线程数: {seg_tpes}, NPU: [1,2]")
    seg_pool = FixedRknnPoolExecutor(
        rknnModel=segmodelPath,
        TPEs=seg_tpes,
        func=seg_func,
        npus=[1, 2]  # 使用NPU 1和2
    )
    
    logger.info(f"初始化姿态RKNN线程池，每NPU线程数: {pose_tpes}, NPU: [0]")
    pose_pool = FixedRknnPoolExecutor(
        rknnModel=posemodelPath,
        TPEs=pose_tpes,
        func=pose_func,
        npus=[0]  # 使用NPU 0
    )
    
    return seg_pool, pose_pool

# 主函数
def main():
    logger.info("初始化RKNN线程池...")
    seg_pool, pose_pool = initialize_pools()
    
    # 创建并启动异步处理器
    logger.info("启动异步处理器，目标帧率: 30 FPS")  # 降低目标帧率便于测试
    async_processor = AsyncProcessor(seg_pool, pose_pool, target_fps=30)
    async_processor.start()
    
    # 结果处理线程
    def result_processing_thread():
        """结果处理线程"""
        while async_processor.running or async_processor.pending_count() > 0:
            result = async_processor.get_result(block=True, timeout=1.0)
            # print("result:",result)
            if result:
                logger.info(f"处理结果: 任务ID={result['task_id']}, "
                           f"推理时间={result['inference_time']*1000:.1f}ms, "
                           f"seg成功={result['seg_success']}, "
                           f"pose成功={result['pose_success']}")
    
    # 启动结果处理线程
    result_thread = threading.Thread(target=result_processing_thread, name="ResultProcessor", daemon=True)
    result_thread.start()
    
    # 主线程监控
    try:
        logger.info("主线程运行中，按Ctrl+C退出...")
        while async_processor.running:
            time.sleep(5)
            # 详细状态监控
            status = (f"任务状态: 待处理={async_processor.pending_count()}, "
                     f"seg线程池待处理={seg_pool.pending_count()}, "
                     f"pose线程池待处理={pose_pool.pending_count()}")
            logger.info(status)
            
    except KeyboardInterrupt:
        logger.info("接收到中断信号，正在停止...")
    finally:
        async_processor.stop()
        logger.info("程序已安全停止")

if __name__ == "__main__":
    main()